package drds.datasource.connection_pool.impl;

import drds.datasource.connection_pool.api.SqlException;
import drds.datasource.connection_pool.impl.holder.ConnectionHolder;
import drds.datasource.connection_pool.util.JdbcUtils;
import drds.datasource.connection_pool.util.MySqlUtils;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;

import java.net.SocketTimeoutException;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.List;

@Slf4j
public class PooledStatement implements drds.datasource.connection_pool.api.Statement {

    @Setter
    @Getter
    private final Statement statement;
    @Setter
    @Getter
    protected PooledConnection pooledConnection;
    @Setter
    @Getter
    protected List<PooledResultSet> pooledResultSetList;
    @Setter
    @Getter
    protected boolean closed = false;

    public PooledStatement(PooledConnection pooledConnection, Statement statement) {
        this.pooledConnection = pooledConnection;
        this.statement = statement;
    }

    //
    @Override
    public final boolean isClosed() throws SqlException {
        return closed;
    }
    //

    @Override
    public final void cancel() throws SqlException {
        checkIsOpen();
        try {
            statement.cancel();
        } catch (SQLException t) {
            throw checkException(t, "statement.cancel()");
        }
    }

    protected void addPooledResultSet(PooledResultSet pooledResultSet) {
        if (pooledResultSetList == null) {
            pooledResultSetList = new ArrayList<PooledResultSet>(1);
        }
        pooledResultSetList.add(pooledResultSet);
    }

    protected void clearPooledResultSetList() throws SqlException {
        if (pooledResultSetList == null) {
            return;
        }
        for (PooledResultSet pooledResultSet : pooledResultSetList) {
            if (pooledResultSet != null) {
                if (!pooledResultSet.isClosed()) {
                    pooledResultSet.close();
                }
            }
        }
        pooledResultSetList.clear();
    }

    //
    @Override
    public final drds.datasource.connection_pool.api.Connection getConnection() throws SqlException {
        checkIsOpen();
        return pooledConnection;
    }

    @Override
    public void close() throws SqlException {
        if (this.closed) {
            return;
        }
        clearPooledResultSetList();
        if (statement != null) {
            try {
                statement.close();
            } catch (SQLException e) {
                JdbcUtils.log(log, "statement.close()", e);
                //
                throw new SqlException(e);
            }
        }
        this.closed = true;
        //
        ConnectionHolder connectionHolder = pooledConnection.getConnectionHolder();
        if (connectionHolder != null) {
            connectionHolder.removePooledStatement(this);
        }
    }


    /**
     * 异常校验
     */
    protected SqlException checkException(SQLException e, String method) throws SqlException {
        JdbcUtils.log(log, method, e);
        //
        String sql = null;
        handleScoketTimeout(e);
        return pooledConnection.handleException(e, method, sql);
    }

    protected SqlException checkException(SQLException e, String method, String sql) throws SqlException {
        JdbcUtils.log(log, method, e);
        //
        handleScoketTimeout(e);
        return pooledConnection.handleException(e, method, sql);
    }

    protected void handleScoketTimeout(Throwable error) throws SqlException {
        if (this.pooledConnection == null || this.pooledConnection.connectionHolder == null) {
            return;
        }
        DataSource dataSource = (DataSource) this.pooledConnection.connectionHolder.dataSource;
        if (dataSource == null) {
            return;
        }
        if (!dataSource.killWhenSocketReadTimeout) {
            return;
        }
        SQLException sqlException = null;
        if (error instanceof SQLException) {
            sqlException = (SQLException) error;
        }
        if (sqlException == null) {
            return;
        }
        Throwable cause = error.getCause();
        boolean socketReadTimeout = cause instanceof SocketTimeoutException && "Read timed out".equals(cause.getMessage());
        if (!socketReadTimeout) {
            return;
        }
        String killQuery = MySqlUtils.buildKillQuerySql(this.pooledConnection.getConnection(), (SQLException) error);
        if (killQuery == null) {
            return;
        }
        drds.datasource.connection_pool.api.Connection connection = null;
        drds.datasource.connection_pool.api.Statement statement = null;
        try {
            connection = dataSource.getConnection(1000);
            if (connection == null) {
                return;
            }
            statement = connection.createStatement();
            statement.executeQuery(killQuery);
            log.info(killQuery + " success.");
        } catch (SqlException ex) {
            log.warn(killQuery + " e.", ex);
        } finally {
            statement.close();
            pooledConnection.close();
        }
    }

    //
    protected void checkIsOpen() throws SqlException {
        if (closed) {
            Throwable disableError = null;
            if (this.pooledConnection != null) {
                disableError = this.pooledConnection.getDisableError();
            }

            if (disableError != null) {
                throw new SqlException("preparedStatement is closed", disableError);
            } else {
                throw new SqlException("preparedStatement is closed");
            }
        }
    }

    //
    @Override
    public final int executeUpdate(String sql) throws SqlException {
        checkIsOpen();
        pooledConnection.beforeExecute();
        try {
            return statement.executeUpdate(sql);
        } catch (SQLException e) {
            throw checkException(e, "statement.executeUpdate(sql)", sql);
        } finally {
            pooledConnection.afterExecute();
        }
    }


    @Override
    public final int executeUpdate(String sql, int autoGeneratedKeys) throws SqlException {
        checkIsOpen();
        pooledConnection.beforeExecute();
        try {
            return statement.executeUpdate(sql, autoGeneratedKeys);
        } catch (SQLException e) {
            throw checkException(e, "statement.executeUpdate(sql, autoGeneratedKeys)", sql);
        } finally {
            pooledConnection.afterExecute();
        }
    }

    @Override
    public final drds.datasource.connection_pool.api.ResultSet getGeneratedKeys() throws SqlException {
        checkIsOpen();
        try {
            ResultSet resultSet = statement.getGeneratedKeys();
            PooledResultSet pooledResultSet = new PooledResultSet(this, resultSet);
            addPooledResultSet(pooledResultSet);
            return pooledResultSet;
        } catch (SQLException e) {
            throw checkException(e, "statement.getGeneratedKeys()");
        }
    }

    @Override
    public final int getUpdateCount() throws SqlException {
        checkIsOpen();
        try {
            return statement.getUpdateCount();
        } catch (SQLException e) {
            throw checkException(e, "statement.getUpdateCount()");
        }
    }

    //
    @Override
    public final void addBatch(String sql) throws SqlException {
        checkIsOpen();
        try {
            statement.addBatch(sql);
        } catch (SQLException e) {
            throw checkException(e, "statement.addBatch(sql)", sql);
        }
    }

    @Override
    public int[] executeBatch() throws SqlException {
        checkIsOpen();
        pooledConnection.beforeExecute();
        try {
            return statement.executeBatch();
        } catch (SQLException e) {
            throw checkException(e, "statement.executeBatch()");
        } finally {
            pooledConnection.afterExecute();
        }
    }

    @Override
    public final void clearBatch() throws SqlException {
        if (closed) {
            return;
        }
        try {
            statement.clearBatch();
        } catch (SQLException e) {
            throw checkException(e, "statement.clearBatch()");
        }
    }
    //

    @Override
    public final drds.datasource.connection_pool.api.ResultSet executeQuery(String sql) throws SqlException {
        checkIsOpen();
        pooledConnection.beforeExecute();
        try {
            ResultSet resultSet = statement.executeQuery(sql);
            if (resultSet == null) {
                return null;
            }
            PooledResultSet pooledResultSet = new PooledResultSet(this, resultSet);
            addPooledResultSet(pooledResultSet);
            return pooledResultSet;
        } catch (SQLException e) {
            throw checkException(e, "statement.executeQuery(sql)", sql);
        } finally {
            pooledConnection.afterExecute();
        }
    }


    @Override
    public final drds.datasource.connection_pool.api.ResultSet getResultSet() throws SqlException {
        checkIsOpen();
        try {
            ResultSet resultSet = statement.getResultSet();
            if (resultSet == null) {
                return null;
            }
            PooledResultSet pooledResultSet = new PooledResultSet(this, resultSet);
            addPooledResultSet(pooledResultSet);
            return pooledResultSet;
        } catch (SQLException e) {
            throw checkException(e, "statement.getResultSet()");
        }
    }
    //

    @Override
    public final int getQueryTimeout() throws SqlException {
        checkIsOpen();
        try {
            return statement.getQueryTimeout();
        } catch (SQLException e) {
            throw checkException(e, "statement.getQueryTimeout()");
        }
    }

    @Override
    public void setQueryTimeout(int seconds) throws SqlException {
        checkIsOpen();
        try {
            statement.setQueryTimeout(seconds);
        } catch (SQLException e) {
            throw checkException(e, "statement.setQueryTimeout(seconds)");
        }
    }

    @Override
    public final int getFetchSize() throws SqlException {
        checkIsOpen();
        try {
            return statement.getFetchSize();
        } catch (SQLException e) {
            throw checkException(e, "statement.getFetchSize()");
        }
    }

    @Override
    public void setFetchSize(int rows) throws SqlException {
        checkIsOpen();
        try {
            statement.setFetchSize(rows);
        } catch (SQLException e) {
            throw checkException(e, "statement.setFetchSize(rows)");
        }
    }


}
